package com.raven.play.stream.test.javadsl.tutorial;

import akka.NotUsed;
import akka.stream.*;
import akka.stream.javadsl.Sink;
import akka.stream.javadsl.Source;
import akka.stream.stage.*;
import com.raven.play.stream.test.BaseTest;
import org.junit.jupiter.api.Test;
import scala.concurrent.Future;
import scala.concurrent.Promise;

import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Random;

public class $23CustomOperators extends BaseTest {

    @Test
    public void test1() {
        Source<Integer, NotUsed> source = Source.fromGraph(new RandomNumberGenerator(100)).throttle(1, Duration.of(1, ChronoUnit.SECONDS));
        Sink<Integer, NotUsed> sink = Sink.fromGraph(new Batcher(10));
        // source.to(Sink.foreach(it -> log.info(it.toString()))).run(system);
        source.to(sink).run(system);
    }

    class RandomNumberGenerator extends GraphStage<SourceShape<Integer>> {

        private final Integer max;
        private final Random random;
        // 组件输出端口
        private final Outlet<Integer> outPort;

        public RandomNumberGenerator(Integer max) {
            this.max = max;
            this.random = new Random();
            this.outPort = Outlet.create("randomGenerator");
        }

        @Override
        public SourceShape<Integer> shape() {
            return new SourceShape<>(outPort);
        }

        @Override
        public GraphStageLogic createLogic(Attributes inheritedAttributes) throws Exception, Exception {
            return new GraphStageLogic(shape()) {{
                // 定义可变状态和实现逻辑
                setHandler(outPort, new OutHandler() {
                    @Override
                    public void onPull() {
                        // 当下游请求时，发出一个新元素
                        int nextNumber = random.nextInt(max);
                        // 通过outPort推送它
                        push(outPort, nextNumber);
                        log.info("Sent: {}", nextNumber);
                    }
                });
            }};
        }
    }

    class Batcher extends GraphStage<SinkShape<Integer>> {

        private final int batchSize;
        private final Inlet<Integer> inPort;

        public Batcher(int batchSize) {
            this.batchSize = batchSize;
            this.inPort = Inlet.create("batcher");
        }

        @Override
        public SinkShape<Integer> shape() {
            return new SinkShape<>(inPort);
        }

        @Override
        public GraphStageLogic createLogic(Attributes inheritedAttributes) throws Exception, Exception {
            return new GraphStageLogic(shape()) {
                // 可变状态
                private final Queue<Integer> batch = new LinkedList<>();

                @Override
                public void preStart() throws Exception, Exception {
                    pull(inPort);
                }

                {
                    setHandler(inPort, new InHandler() {
                        @Override
                        public void onPush() {
                            // 当上游要发送元素时
                            int nextElement = grab(inPort);
                            batch.add(nextElement);
                            log.info("Received: {}", nextElement);

                            // 假设一些复杂的计算
                            try {
                                Thread.sleep(100);
                            } catch (InterruptedException e) {
                                // 处理中断异常
                            }

                            if (batch.size() >= batchSize) {
                                System.out.println("New batch: " + batch);
                                batch.clear();
                            }

                            // 向上游发送需求
                            pull(inPort);
                        }

                        @Override
                        public void onUpstreamFinish() {
                            // 当上游完成时
                            if (!batch.isEmpty()) {
                                System.out.println("New batch: " + batch);
                                System.out.println("Stream finished.");
                                batch.clear();
                            }
                        }

                    });
                }
            };
        }
    }

    class CounterFlow<T> extends GraphStageWithMaterializedValue<FlowShape<T, T>, Future<Integer>> {
        private final Inlet<T> inPort = Inlet.create("counterIn");
        private final Outlet<T> outPort = Outlet.create("counterOut");

        @Override
        public FlowShape<T, T> shape() {
            return new FlowShape<>(inPort, outPort);
        }

        @Override
        public scala.Tuple2<GraphStageLogic, Future<Integer>> createLogicAndMaterializedValue(Attributes inheritedAttributes) {
            Promise<Integer> promise = scala.concurrent.Promise.apply();

            GraphStageLogic logic = new GraphStageLogic(shape()) {
                private int counter = 0;

                {
                    setHandler(outPort, new OutHandler() {
                        @Override
                        public void onPull() {
                            pull(inPort);
                        }

                        @Override
                        public void onDownstreamFinish() throws Exception{
                            promise.success(counter);
                            OutHandler.super.onDownstreamFinish();
                        }
                    });

                    setHandler(inPort, new InHandler() {
                        @Override
                        public void onPush() {
                            T nextElement = grab(inPort);
                            counter++;
                            push(outPort, nextElement);
                        }

                        @Override
                        public void onUpstreamFinish() throws Exception{
                            promise.success(counter);
                            InHandler.super.onUpstreamFinish();
                        }

                        @Override
                        public void onUpstreamFailure(Throwable ex) throws Exception {
                            promise.failure(ex);
                            InHandler.super.onUpstreamFailure(ex);
                        }
                    });
                }
            };

            return new scala.Tuple2<>(logic, promise.future());
        }
    }

}
